[MXNET-1382] Add the index_array operator #14638

nickguletskii · 2019-04-07T16:24:24Z

Description

This pull request implements index_array, an operator that returns an array of indexes of the input array.

For an input array with shape (d_1, d_2, ..., d_n), index_array returns a (d_1, d_2, ..., d_n, n) array idx, where idx[i_1, i_2, ..., i_n, :] = [i_1, i_2, ..., i_n].

Additionally, when the parameter axes is specified, idx will be a (d_1, d_2, ..., d_n, m) array where m is the length of axes, and the following
equality will hold: idx[i_1, i_2, ..., i_n, j] = i_{axes[j]}.

Examples:

x = mx.nd.ones((3, 2))

mx.nd.contrib.index_array(x) = [[[0 0]
                                 [0 1]]

                                [[1 0]
                                 [1 1]]

                                [[2 0]
                                 [2 1]]]

x = mx.nd.ones((3, 2, 2))

mx.nd.contrib.index_array(x, axes=(1, 0)) = [[[[0 0]
                                               [0 0]]

                                              [[1 0]
                                               [1 0]]]


                                             [[[0 1]
                                               [0 1]]

                                              [[1 1]
                                               [1 1]]]

Motivation

This operator can be used to generate meshgrids for tensors without knowing their exact shapes during construction. For instance, this operator can be used to make a makeshift prior box generator for anchor-based computer vision models:

feature_map = F.ones((8, 128, 128, 256)) # N x H x W x C, no shape information when using the Symbol API.
prior_box_stride = 16
box_size=[8, 8]

template = F.squeeze(F.slice_axis(feature_map, begin=0, end=1, axis=-1), axis=-1) # N x H x W
box_centres = F.contrib.index_array(template, axes=(-2, -1, -2, -1)).astype("float32") # N x H x W x 4
box_centres = F.broadcast_mul(box_centres, F.array([prior_box_stride]).reshape((1, 1, 1, 1))) # N x H x W x 4
corner_offsets = F.array(box_size).reshape((1, 1, 1, 2))
corner_offsets = F.concat(-corner_offsets/2, corner_offsets/2, dim=-1)
box_corners = F.broadcast_plus(box_centres, corner_offsets)

Also, this operator can be applied to implement positional encodings for sequence processing, e.g.:

sequence_embeddings = F.ones((65, 8, 256)) # T x N x C, no shape information when using the Symbol API.
template = sequence_embeddings.reshape((0, 0, -1, 2)) # T x N x C -> T x N x (C/2) x 2
pos, i = F.split(
    F.contrib.index_array(template, axes=(0, 2)).astype("float32"), # T x N x (C/2) x 2 x 2
    axis=-1,
    num_outputs=2,
    squeeze_axis=True
) # T x N x (C/2) x 2 and T x N x (C/2) x 2
base = F.ones((1, 1, 1, 1)) * 10000
dmodel = F.slice_axis(F.shape_array(sequence_embeddings), begin=-1, end=None, axis=0)
dmodel = dmodel.reshape((1, 1, 1, 1)).astype("float32")
tmp = F.broadcast_div(pos, F.broadcast_power(base, F.broadcast_div(2 * i,  dmodel))) # T x N x (C/2) x 2
sin_input, cos_input = F.split(tmp, axis=-1, num_outputs=2, squeeze_axis=True) # T x N x (C/2) and T x N x (C/2)
positional_encoding = F.stack(F.sin(sin_input), F.cos(cos_input), axis=-1).reshape((0, 0, -3)) # T x N x C

I've also encountered situations where this operator would have been useful for some indexing tricks.

Checklist

Essentials

Please feel free to remove inapplicable items for your PR.

The PR title starts with [MXNET-$JIRA_ID], where $JIRA_ID refers to the relevant JIRA issue created (except PRs with tiny changes)
Changes are complete (i.e. I finished coding on this PR)
All changes have test coverage:
Unit tests are added for small changes to verify correctness (e.g. adding a new operator)
Nightly tests are added for complicated/long-running ones (e.g. changing distributed kvstore)
Build tests will be added for build configuration changes (e.g. adding a new build option with NCCL)
Code is well-documented:
For user-facing API changes, API doc string has been updated.
For new C++ functions in header files, their functionalities and arguments are documented.
For new examples, README.md is added to explain the what the example does, the source of the dataset, expected performance on test set and reference to the original paper if applicable
Check the API doc at http://mxnet-ci-doc.s3-accelerate.dualstack.amazonaws.com/PR-$PR_ID/$BUILD_ID/index.html
To the my best knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change

Changes

The IndexArray operator (for the CPU)
The IndexArray operator (for the GPU)
Tests for the CPU implementation
Tests for the GPU implementation
Entries on the Python API reference pages

Comments

This operator always returns kInt64.
The new tests in test_operator_gpu.py are exactly the same as in test_operator.py. I couldn't find any evidence that test_operator.py ever gets called with a GPU default_context, so I copied the tests into test_operator_gpu.py to make sure that the GPU implementation works too.

piyushghai · 2019-04-07T18:28:52Z

Thanks for your contributions @nickguletskii. Can you also look into the CI failures ?
@mxnet-label-bot Add [Opeartor, pr-awaiting-review]

piyushghai · 2019-04-07T18:29:15Z

@mxnet-label-bot Update [Operator, pr-awaiting-review]

nickguletskii · 2019-04-07T19:13:14Z

The CI failure seems to be unrelated. Here's the relevant part of the log:

[2019-04-03 23:57:30   ERROR] Cuda initialization failure with error 35. Please check cuda installation:  http://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html.

terminate called after throwing an instance of 'std::runtime_error'

  what():  Failed to create object

If I am not mistaken, this error code corresponds to the following error message: "CUDA driver version is insufficient for CUDA runtime version". Could this be caused by the recent update to the nvidia/cuda:10.0-cudnn7-devel Docker tag? The TensorRT job uses nvidia/cuda:10.0-cudnn7-devel, while other jobs use nvidia/cuda:9.1-cudnn7-devel (e.g. the ubuntu_gpu build), which may be why the other jobs are not failing.

fhieber · 2019-04-09T15:54:39Z

Thanks for implementing this! This seems similar to what I requested in #14574, correct?

Any reason this should be in the contrib namespace?

nickguletskii · 2019-04-09T21:39:57Z

@fhieber In a way, yes. The behaviour you have described in #14574 can be emulated by applying reshape, index_array, broadcast_add, and finally, reshape_list. I was thinking of implementing arange_like instead of index_array, but I don't find the arange operator itself versatile enough and I didn't want to make arange_like too different from arange.

There's no particular reason why I've placed index_array into the contrib namespace other than the fact that it's a new operator. I'd be happy to update this pull request and move index_array out of the contrib namespace if the community deems it desirable.

fhieber · 2019-04-11T09:17:40Z

Thanks @nickguletskii , I agree, your operator seems more versatile, thanks for working on it!

contrib is fine with me, I just asked out of curiosity on whether the preferred or standard way of introducing new operators is to always place in contrib first.

nickguletskii · 2019-04-12T16:57:21Z

@piyushghai Sorry to bother you, but should I rebase this pull request on master now that the failing job (test_unix_python3_tensorrt_gpu) has been disabled (#14642), or should I just leave it as it is and wait for the review?

piyushghai · 2019-04-12T16:59:11Z

@nickguletskii You should rebase with the latest master branch. That'll enable you to push again, allowing the CI to pass on your PR.
The reviews will happen on it in the meanwhile :)

nickguletskii · 2019-04-30T23:20:03Z

Huh, this is weird, I've updated this PR to be compatible with the recently merged support for zero-dim and zero-shape arrays, and now the CI for Windows CPU tests are failing. However, when I try to run the tests locally (on a 64 bit Windows machine), the tests pass just fine. I'll make further attempts to reproduce this issue on my local machine tomorrow. I suspect that I've screwed up with memory access somewhere because the test runner crashes without providing any useful debugging information.

nickguletskii · 2019-05-02T07:19:00Z

The previous CI failures on Windows were caused by indexing C-style arrays using ints. I'm not sure whether this was just undefined behaviour or a bug in the older versions of MSVC++, but it did not happen with VC++ 16.0. I fixed the issue by using ptrdiff_t to index C-style arrays.

nickguletskii · 2019-05-02T07:23:37Z

@szha @reminisce This PR should be ready for a review now. The CI failure seems to be unrelated (the failing test is test_operator_gpu.test_convolution_multiple_streams, described in #14329).

vandanavk · 2019-05-10T05:41:40Z

@szha @reminisce @fhieber Is this PR good to go?

reminisce

Thank you for your contribution and sorry for the late review. Have some comments on code cosmetics.

src/operator/contrib/index_array.cc

src/operator/contrib/index_array-inl.h

reminisce · 2019-05-13T06:22:15Z

tests/python/gpu/test_operator_gpu.py

@@ -2130,6 +2130,68 @@ def test_bilinear_sampler_versions():
                if req_dict['grid'] is 'write':
                    assert_almost_equal(exe.grad_dict['grid'].asnumpy(), exe_list[ref_idx].grad_dict['grid'].asnumpy(), rtol=1e-3, atol=1e-5)

+@with_seed()


The tests defined in test_operator.py will be run with ctx=mx.gpu() in CI. There is no need to duplicate test code in this file.

Sorry, I've removed these redundant tests in 3cdfcf5.

tests/python/unittest/test_operator.py

reminisce · 2019-05-13T06:27:50Z

tests/python/unittest/test_operator.py

+
+@with_seed()
+def test_index_array_default_zero_dim():
+    with mx.np_compat(active=True):


If the whole function is expected to be numpy compatible, you can use a decorator mx.use_np_compat to wrap the whole function. It would be easier for us to just delete the decorator in the next major release than putting additional efforts of adjusting indentation.

Thanks, I've updated the subtests to use @mx.use_np_compat in 139567b.

reminisce · 2019-05-13T06:28:18Z

tests/python/unittest/test_operator.py

+
+@with_seed()
+def test_index_array_default_zero_size():
+    with mx.np_compat(active=True):


Same as above, please use the decorator mx.use_np_compat.

reminisce · 2019-05-13T06:28:47Z

tests/python/unittest/test_operator.py

+
+@with_seed()
+def test_index_array_select_axes_zero_size():
+    with mx.np_compat(active=True):


Same as above.

karan6181 · 2019-05-21T22:17:20Z

@nickguletskii Could you please re-trigger the CI again?
@szha Could you please review this PR? Thanks!

Changes the implementation of index_array to be compatible with the recently merged support for zero-dim and zero-size arrays. Resolves the incompatibilities with apache#14661.

In the previous implementation, the output gradient had an incorrect shape. This commit fixes the shapes and makes the tests more readable.

Solves access violations when compiling with MSVC++ 14.0.

* Implement the index_array operator * Add index_array operator tests * Add index_array operator GPU tests * Add the index_array operator to the Python docs autosummary * Add the author of the index_array operator to CONTRIBUTORS.md * Make index_array compatible with zero-dim and zero-size arrays Changes the implementation of index_array to be compatible with the recently merged support for zero-dim and zero-size arrays. Resolves the incompatibilities with apache#14661. * Fix the index_array gradient checks in the unit tests In the previous implementation, the output gradient had an incorrect shape. This commit fixes the shapes and makes the tests more readable. * Add zero-dim and zero-size array tests for index_array * Use mxnet::Tuple<int> instead of TShape for the axes parameter * Fix incorrect array indexing in index_array Solves access violations when compiling with MSVC++ 14.0. * Avoid copying the input shape array in the index_array shape function * Add unknown shape handling to index_array * Use SHAPE_ASSIGN_CHECK to assign the shape in index_array * Remove the redundant index_array GPU tests from test_operator_gpu.py * Move the index_array tests into a single function (test_index_array) * Use @mx.use_np_compat instead of mx.np_compat in index_array op tests * Remove the use of template specialization for IndexArrayForward * Add the index_array operator to the AMP symbol list * Retrigger CI

nickguletskii requested a review from szha as a code owner April 7, 2019 16:24

marcoabreu added the pr-awaiting-review PR is waiting for code review label Apr 7, 2019

marcoabreu added the Operator label Apr 7, 2019

nickguletskii force-pushed the feature/index-array-op branch from 4ea3cf6 to 0f742af Compare April 12, 2019 17:00

nickguletskii force-pushed the feature/index-array-op branch from 7fa5a17 to aba94af Compare April 30, 2019 18:42

szha requested a review from reminisce April 30, 2019 21:22

reminisce reviewed May 13, 2019

View reviewed changes

reminisce approved these changes May 14, 2019

View reviewed changes

nickguletskii added 9 commits May 22, 2019 19:01

Implement the index_array operator

a9451cb

Add index_array operator tests

91e2052

Add index_array operator GPU tests

1a4a0b6

Add the index_array operator to the Python docs autosummary

c58964a

Add the author of the index_array operator to CONTRIBUTORS.md

50c0c28

Make index_array compatible with zero-dim and zero-size arrays

df1f421

Changes the implementation of index_array to be compatible with the recently merged support for zero-dim and zero-size arrays. Resolves the incompatibilities with apache#14661.

Fix the index_array gradient checks in the unit tests

e4152c6

In the previous implementation, the output gradient had an incorrect shape. This commit fixes the shapes and makes the tests more readable.

Add zero-dim and zero-size array tests for index_array

adb9f36

Use mxnet::Tuple<int> instead of TShape for the axes parameter

3eccb3c

nickguletskii added 9 commits May 22, 2019 19:01

Fix incorrect array indexing in index_array

a6b3cab

Solves access violations when compiling with MSVC++ 14.0.

Avoid copying the input shape array in the index_array shape function

69bf4f6

Add unknown shape handling to index_array

be39798

Use SHAPE_ASSIGN_CHECK to assign the shape in index_array

445f758

Remove the redundant index_array GPU tests from test_operator_gpu.py

972b05f

Move the index_array tests into a single function (test_index_array)

3051d5e

Use @mx.use_np_compat instead of mx.np_compat in index_array op tests

3f8702a

Remove the use of template specialization for IndexArrayForward

496bad3

Add the index_array operator to the AMP symbol list

69a9969

nickguletskii force-pushed the feature/index-array-op branch from 89e9cf5 to 69a9969 Compare May 22, 2019 16:06

Retrigger CI

ffde9b3

szha merged commit f689042 into apache:master May 25, 2019

nickguletskii deleted the feature/index-array-op branch May 27, 2019 16:20

RogerChern mentioned this pull request Aug 13, 2019

Future Roadmap Suggestions tusen-ai/simpledet#66

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[MXNET-1382] Add the index_array operator #14638

[MXNET-1382] Add the index_array operator #14638

nickguletskii commented Apr 7, 2019 •

edited

Loading

piyushghai commented Apr 7, 2019

piyushghai commented Apr 7, 2019

nickguletskii commented Apr 7, 2019 •

edited

Loading

fhieber commented Apr 9, 2019

nickguletskii commented Apr 9, 2019

fhieber commented Apr 11, 2019

nickguletskii commented Apr 12, 2019

piyushghai commented Apr 12, 2019

nickguletskii commented Apr 30, 2019

nickguletskii commented May 2, 2019

nickguletskii commented May 2, 2019 •

edited

Loading

vandanavk commented May 10, 2019

reminisce left a comment

reminisce May 13, 2019

nickguletskii May 13, 2019

reminisce May 13, 2019

nickguletskii May 13, 2019

reminisce May 13, 2019

reminisce May 13, 2019

karan6181 commented May 21, 2019

[MXNET-1382] Add the index_array operator #14638

[MXNET-1382] Add the index_array operator #14638

Conversation

nickguletskii commented Apr 7, 2019 • edited Loading

Description

Examples:

Motivation

Checklist

Essentials

Changes

Comments

piyushghai commented Apr 7, 2019

piyushghai commented Apr 7, 2019

nickguletskii commented Apr 7, 2019 • edited Loading

fhieber commented Apr 9, 2019

nickguletskii commented Apr 9, 2019

fhieber commented Apr 11, 2019

nickguletskii commented Apr 12, 2019

piyushghai commented Apr 12, 2019

nickguletskii commented Apr 30, 2019

nickguletskii commented May 2, 2019

nickguletskii commented May 2, 2019 • edited Loading

vandanavk commented May 10, 2019

reminisce left a comment

Choose a reason for hiding this comment

reminisce May 13, 2019

Choose a reason for hiding this comment

nickguletskii May 13, 2019

Choose a reason for hiding this comment

reminisce May 13, 2019

Choose a reason for hiding this comment

nickguletskii May 13, 2019

Choose a reason for hiding this comment

reminisce May 13, 2019

Choose a reason for hiding this comment

reminisce May 13, 2019

Choose a reason for hiding this comment

karan6181 commented May 21, 2019

nickguletskii commented Apr 7, 2019 •

edited

Loading

nickguletskii commented Apr 7, 2019 •

edited

Loading

nickguletskii commented May 2, 2019 •

edited

Loading